import os
import argparse


parser = argparse.ArgumentParser(description='modelscope模型下载')

parser.add_argument('--model_name', type=str, help='the model name from modelscope, example AI-ModelScope/stable-diffusion-2-1', required=True)
parser.add_argument('--local_dir', type=str, help='the model cache path.', default=os.getcwd(), required=True)


if __name__ == '__main__':
    args = parser.parse_args()

    print(f"current workspace is {os.getcwd()}")
    print(f"the model_name is {args.local_dir}/{args.model_name}")
    print(f"the local_dir is {args.local_dir}")

    try:
        from modelscope import snapshot_download
        model_dir = snapshot_download(args.model_name, local_dir=args.local_dir)
    except ImportError:
        print("modelscope was not installed! try to install...")
        os.system("pip install modelscope")
    except Exception as e:
        print(f"An error occurred: {e}")
